refactor(rollout): move the driver-side async engines to rollout/manager/ - #7
Closed
celve wants to merge 1 commit into
Closed
refactor(rollout): move the driver-side async engines to rollout/manager/#7celve wants to merge 1 commit into
celve wants to merge 1 commit into
Conversation
…ger/
They hold no model and cannot generate, so calling them engines made 'engine'
mean three things at once: the ABC, a concrete backend, and these. They are the
layer that owns admission, acceptance and disposal over time — a manager.
engine/asynchronous.py -> manager/{buffers,batch,agentic}.py
AsyncBatchRolloutEngine -> BatchManager
AsyncAgenticRolloutEngine -> AgenticManager
'manager' rather than 'scheduler' because this tree already calls three unrelated
things a scheduler: the LR scheduler (utils/scheduler_utils.py), the diffusion
noise scheduler, and SGLang's own scheduler subprocesses.
Adds manager/protocol.py recording the consumer surface both implementations
share, and splits the mechanisms (VersionedBuffer, PendingGroups, root_of) out of
the file that used to hold both engines and their machinery. launch_ceiling stays
with BatchManager, whose two trainers are its only callers, so its ownership note
remains accurate.
Pure move + rename; no behaviour change.
Test plan: an AST comparison against the pre-move file proves VersionedBuffer,
PendingGroups, InflightPool, _InflightJob, root_of and BatchManager byte-identical
modulo the rename. compileall, check-recipe-targets (2434 paths) and
check-experimental-boundaries all pass.
Owner
Author
|
Closing: opened against the fork by mistake. Re-opening against Tencent-Hunyuan/UniRL, with the @addressed primitive and the manager move combined into one PR. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
AsyncBatchRolloutEngineandAsyncAgenticRolloutEnginehold no model and cannotgenerate. Calling them engines made "engine" mean three different things at once:
the
BaseRolloutEngineABC, a concrete backend, and these driver-side buffermanagers. They are the layer that owns admission, acceptance and disposal over
time — a manager.
"manager" rather than "scheduler" because this tree already calls three unrelated
things a scheduler: the LR scheduler (
utils/scheduler_utils.py,LrSchedulerConfig), the diffusion noise scheduler (flow_match_sde_scheduler.py,patch_scheduler.py,patch_wan_scheduler.py), and SGLang's own schedulersubprocesses (
sglang/config.py:178). A fourth meaning would be worse than thecurrent name.
Adds
manager/protocol.pyrecording the consumer surface both implementationsshare, and splits the policy-free mechanisms (
VersionedBuffer,PendingGroups,root_of) out of the file that held both engines and all their machinery.launch_ceilingstays withBatchManager, whose two trainers are its onlycallers, so its ownership note ("trainer-side POLICY... hosted here only because
it is the two batch trainers' one shared torch-free home") remains accurate.
Pure move + rename. No behaviour change.
Related Issue
LIN-693 — first commit of the change that moves the agentic rollout task queue
from rank 0 of the rollout slab to the driver.
Test Plan
engine/asynchronous.pyprovingVersionedBuffer,PendingGroups,InflightPool,_InflightJob,root_ofand
BatchManagerbyte-identical modulo the rename — all six report IDENTICAL.python -m compileall unirl/— passes.python3 lint/check_recipe_targets.py— 2434 recipe_target_paths resolve.python3 lint/check_experimental_boundaries.py— ok.grepsweep for residualengine.asynchronous/ old class names outside thenew package: only intentional historical notes in the moved docstrings.
Not run; reason: no Ray, torch or GPU in the authoring environment, and this
repo's CI is lint-only, so nothing here has been executed. The rename is proven
structurally rather than behaviourally.
Compatibility / Risk
Import paths change for four trainers (
async_ar,async_diffusion,agentic_async,agentic_partial), all updated in this PR. No recipe, config,checkpoint or data-format change; no
_target_in any recipe points at the movedmodule (confirmed by
check_recipe_targets).Anything on a branch importing
unirl.rollout.engine.asynchronouswill need theone-line path update.
Reviewer Notes
Stacked on #6 (
LIN-693/addressed-calls), which is its base — review thatfirst. This PR's diff against that base is the move alone.
Draft, because this is the first of five commits in the intended change. Still
to come on this branch: exposing
run_trajectory/set_stoppingonAgenticRolloutEngine, givingAgenticManagerthe task queue and per-replicaplacement, rewiring all three agentic trainers, and deleting the rank-0
coordinator. The move stands alone and is mergeable on its own merit, but merging
it in isolation buys only the naming.
The design and staged plan are written up outside this repo; the short version is
in the linked issue.
Prepared with AI assistance. The AST proof is what carries the "pure move" claim —
worth re-running rather than eyeballing the diff. Checked for overlapping open
PRs — there are none.
Checklist